import numpy as np
import pyopencl as cl
from openclinfo import *

#mako
N_TYPE = 1
W_P_PAR = 4   #
p_par = np.zeros((N_TYPE,N_TYPE,W_P_PAR),'float32')

#r_cut = 2.5
r_cut = 2.5
r_skin = 0.3

#neighborlist index: 0-n:normal, -1:null (end of)
#in cl kernel, judge -1 is faster than checking j<n_neigh or if? (can use const loop-size and unrolling)
#bond,angle,exclude _lists are static, can be done once at the beginning by CPU

class Neighbor:
    #def __init__(self,name,cl_file,nei_par,contex,flags='-O3'):
    def __init__(self,name=None,cl_file=None,nei_par=None,context=None,flags=''):
        self.name = name #'pair_neighbor_list'
        self.cl_file = cl_file #'neighbor_list.cl'
        #
        r_cut_tot = r_cut + r_skin
        self.r2_cut = r_cut_tot*r_cut_tot 
        #
        cl_code = file(self.cl_file,"r").read()
        self.prg = cl.Program(context,cl_code).build(options=flags)
        #
    #def create(self,que,sglb,sloc,cfg): #allocate and update

    def sort_host(self,que,cfg):
        ngrids = NGX*NGY*NGZ
        sr = np.ndarray((cfg.n_atoms,4),'float32')
        idx = np.ndarray(cfg.n_atoms,'int32')
        grid = self.grid.reshape(ngrids,MAX_PER_GRID)
        n_grid = self.n_grid.reshape(ngrids)
        id = 0
        for i in range(ngrids):
            ni = n_grid[i]
            for j in range(ni):
                iatm = grid[i,j]
                sr[id] = cfg.r_t[iatm]
                idx[iatm] = id
                id += 1
        self.r_t = sr
        self._r_t = cl.Buffer(que.context,cl.mem_flags.READ_WRITE,size=sr.nbytes)
        cl.enqueue_copy(que,self._r_t,self.r_t)
        print 'sorted-ids-1:',idx[:cfg.n_atoms/2]
        print 'sorted-ids-2:',idx[cfg.n_atoms/2:]
        
        for i in range(cfg.n_atoms):
            ni = self.n_nei[i]
            for j in range(ni):
                iold = self.nei[i,j]
                self.nei[i,j] = idx[iold]
        cl.enqueue_copy(que,self._nei,self.nei)
                



    def create_grid_nei(self,que,cfg): #allocate 
#test:
        self.rr_g = np.zeros((cfg.n_atoms,4),'float32')
        self.grid = np.ones((NGX,NGY,NGZ,MAX_PER_GRID),'int32') * (-1) #init all by -1
        self.n_grid = np.zeros((NGX,NGY,NGZ),'int32')
        self.overflow = np.zeros(1,'int32')
        mf = cl.mem_flags
        nbytes = self.rr_g.nbytes
        self._rr_g = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        nbytes = self.grid.nbytes
        self._grid = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        nbytes = self.n_grid.nbytes
        self._n_grid = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        nbytes = 4 #float = 4 bytes
        self._overflow = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
#
        self.nei = np.ones((MAX_NEIGH,cfg.n_atoms),'int32') * (-1) #init all by -1
        self.n_nei = np.zeros(cfg.n_atoms,'int32')
        self.overflow = np.zeros(1,'int32')
        mf = cl.mem_flags
        nbytes = self.nei.nbytes #float = 4 bytes
        self._nei = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        nbytes = self.n_nei.nbytes #float = 4 bytes
        self._n_nei = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        nbytes = 4 #float = 4 bytes
        self._overflow = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        cl.enqueue_copy(que,self._grid,self.grid)
        cl.enqueue_copy(que,self._n_grid,self.n_grid)
        cl.enqueue_copy(que,self._overflow,self.overflow)
        #
        sglb = (GLOBAL_SIZE,)
        sloc = (LOCAL_SIZE,)
        self.prg.create_grid(
                que,sglb,sloc,
                np.int32(cfg.n_atoms),
                cfg.box_lo, cfg.box_hi,
                cfg._r_t, self._rr_g,
                self._grid, self._n_grid,
                self._overflow)
        cl.enqueue_copy(que,self.overflow,self._overflow)
        print 'cal grid overflow:',self.overflow
        cl.enqueue_copy(que,self.grid,self._grid)
        cl.enqueue_copy(que,self.n_grid,self._n_grid)
        #reset overflow to zero
        self.overflow[0] = 0
        cl.enqueue_copy(que,self._overflow,self.overflow)
        self.prg.create_neigh(
                que,sglb,sloc,
                np.int32(cfg.n_atoms),
                np.float32(self.r2_cut),
                cfg.box_l, cfg.box_p,
                cfg._r_t, self._rr_g,
                self._grid, self._n_grid,
                self._nei, self._n_nei,
                self._overflow)
        #sort
        self.prg.sort_inside_neigh(
                que,sglb,sloc,
                np.int32(cfg.n_atoms),
                self._nei, self._n_nei)

        cl.enqueue_copy(que,self.overflow,self._overflow)
        print 'cal neigh overflow:',self.overflow
        cl.enqueue_copy(que,self.nei,self._nei)
        cl.enqueue_copy(que,self.n_nei,self._n_nei)


#test:
    def create_test(self,que,cfg): #allocate and update
        self.gen_test(que.context,cfg.n_atoms)

    def create_slow(self,que,cfg): #allocate and update
#test:
        #self.nei = np.ones((cfg.n_atoms,MAX_NEIGH),np.int32) * (-1) #init all by -1
        self.nei = np.ndarray((cfg.n_atoms,MAX_NEIGH),'int32') #init all by -1
        #self.nei = self.nei.astype(np.int32)
        self.n_nei = np.zeros(cfg.n_atoms,'int32')
        self.overflow = np.zeros(1,'int32')
        mf = cl.mem_flags
        nbytes = cfg.n_atoms * 4*MAX_NEIGH #float = 4 bytes
        print 'nbytes:',nbytes, self.nei.nbytes
        self._nei = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        nbytes = cfg.n_atoms * 4 #float = 4 bytes
        self._n_nei = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        nbytes = 4 #float = 4 bytes
        self._overflow = cl.Buffer(que.context,mf.READ_WRITE,size=nbytes)
        cl.enqueue_copy(que,self._overflow,self.overflow)
        #
        sglb = (GLOBAL_SIZE,)
        sloc = (LOCAL_SIZE,)
        self.prg.create_neigh_slow(
                que,sglb,sloc,
                np.int32(cfg.n_atoms),
                np.float32(self.r2_cut),
                cfg.box_l, cfg.box_p,
                cfg._r_t,
                self._nei, self._n_nei,
                self._overflow)
#
        cl.enqueue_copy(que,self.overflow,self._overflow)
        print 'cal neigh overflow:',self.overflow
        cl.enqueue_copy(que,self.nei,self._nei)
        cl.enqueue_copy(que,self.n_nei,self._n_nei)

    def update(self): #only update
        pass

#test:
    def gen_test(self,ctx,n):
        neigh = np.ones((n,MAX_NEIGH),'int32') * (-1) #init all by -1
        n_neigh = np.zeros(n,'int32')
        for i in range(n):
            k = int(np.random.rand()*MAX_NEIGH)
            n_neigh[i] = k
            neigh[i,0:k] = np.random.rand(k)*n
        #print 'neigh-0:',neigh[0]
        self.nei = neigh
        self.n_nei = n_neigh
        print 'max-id:',np.max(neigh)
        mf = cl.mem_flags
        self._nei = cl.Buffer(ctx,mf.READ_ONLY|mf.COPY_HOST_PTR,hostbuf=neigh)
        self._n_nei = cl.Buffer(ctx,mf.READ_ONLY|mf.COPY_HOST_PTR,hostbuf=n_neigh)


